package com.itheima.ai.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.itheima.ai.entity.po.SpringAiChatMemory;
import com.itheima.ai.mapper.SpringAiChatMemoryMapper;
import com.itheima.ai.service.ISpringAiChatMemoryService;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.lang.NonNull;
import org.springframework.stereotype.Service;

import java.time.LocalDateTime;
import java.util.List;
import java.util.Optional;

@Service
public class SpringAiChatMemoryServiceImpl
        extends ServiceImpl<SpringAiChatMemoryMapper, SpringAiChatMemory>
        implements ISpringAiChatMemoryService{

    @NonNull
    @Override
    public List<String> findConversationIds() {
        return Optional.ofNullable(this.getBaseMapper().listConversationIds()).orElse(List.of());
    }

    @NonNull
    @Override
    public List<Message> findByConversationId(@NonNull String conversationId) {
        // SELECT content, type FROM SPRING_AI_CHAT_MEMORY WHERE conversation_id = ? ORDER BY `timestamp`
        List<SpringAiChatMemory> memoryList = this.lambdaQuery()
                .select(SpringAiChatMemory::getContent, SpringAiChatMemory::getType)
                .eq(SpringAiChatMemory::getConversationId, conversationId)
                .orderByAsc(SpringAiChatMemory::getTimestamp)
                .list(); // 查询前20条
        return Optional.ofNullable(memoryList)
                .orElse(List.of())
                .stream()
                .map(this::toMessage)
                .toList();
    }

    private Message toMessage(SpringAiChatMemory memory) {
        return memory.getType().equals(MessageType.USER.name()) ? new UserMessage(memory.getContent()) : new AssistantMessage(memory.getContent());
    }

    @Override
    public void saveAll(@NonNull String conversationId, List<Message> messages) {
        List<SpringAiChatMemory> list = messages.stream()
                .map(message -> toMemory(conversationId, message))
                .toList();
        this.saveBatch(list);
    }

    private SpringAiChatMemory toMemory(String conversationId, Message message) {
        SpringAiChatMemory memory = new SpringAiChatMemory();
        memory.setType(message.getMessageType().name());
        memory.setContent(message.getText());
        memory.setConversationId(conversationId);
        memory.setTimestamp(LocalDateTime.now());
        return memory;
    }

    @Override
    public void deleteByConversationId(@NonNull String conversationId) {
        remove(new LambdaQueryWrapper<SpringAiChatMemory>().eq(SpringAiChatMemory::getConversationId, conversationId));
    }

    @Override
    public List<Message> findTopN(@NonNull String conversationId, int size) {
        List<SpringAiChatMemory> memoryList = this.lambdaQuery()
                .select(SpringAiChatMemory::getContent, SpringAiChatMemory::getType)
                .eq(SpringAiChatMemory::getConversationId, conversationId)
                .orderByAsc(SpringAiChatMemory::getTimestamp)
                .list(Page.of(0, size)); // 查询前N
        return Optional.ofNullable(memoryList)
                .orElse(List.of())
                .stream()
                .map(this::toMessage)
                .toList();
    }
}